use core::{
    borrow::BorrowMut,
    ffi::c_void,
    intrinsics::{likely, unlikely},
    mem::{size_of, ManuallyDrop},
    ops::{Deref, DerefMut},
    sync::atomic::{AtomicI32, AtomicU32},
};

use bitflags::bitflags;
use cpuindex::this_processor_id;
use irqflags::{local_irq_restore, local_irq_save};
use raw_list::{RawListHead, RawListNode};

use crate::{
    addrspace::PfnAddr,
    error::Result,
    global::zone_idx,
    zone::{FreeArea, PerdataPages, ZoneInner, MAX_ORDER},
    GfpFlags,
};

#[inline(always)]
fn prefetch(page: *const Page) {
    unsafe {
        #[cfg(target_arch = "aarch64")]
        core::arch::aarch64::_prefetch::<
            { core::arch::aarch64::_PREFETCH_READ },
            { core::arch::aarch64::_PREFETCH_LOCALITY3 },
        >(page as *const i8);
    }
}

bitflags! {
    #[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
    struct PageFlags: u32 {
        // 页属于伙伴系统
        const PageBuddy       = 0b0000000001;
        // 页属于 slab 系统
        const PageSlab        = 0b0000000010;

        // 页属于 Dma zone 区域
        const ZoneDma         = 0b0000000100;
        // 页属于 Normal zone 区域
        const ZoneNormal      = 0b0000001000;

        // 页属于预留页, 不会被使用
        const PageReserved    = 0b0000010000;

        // 页被初始化, 首次对页进行初始化
        const PageInitialized = 0b0000100000;

        // 当页是复合页时, 标记页是头部页
        const PageHead        = 0b0001000000;
        // 当页是复合页时, 标记页是尾部页(尾随部分)
        const PageTail        = 0b0010000000;

        // zone 区掩码
        const ZoneMask        = PageFlags::ZoneDma.bits() | PageFlags::ZoneNormal.bits();

        // 当一个页释放时检查使用. 页在释放前不应该包含这些 flags. 如果有, 则有问题.
        const CheckAtFree     = PageFlags::PageSlab.bits() | PageFlags::PageReserved.bits();
        // 当一个页预备返回 buddy 时, 清除这些标志
        const CheckAtPrep     = PageFlags::PageBuddy.bits() | PageFlags::PageSlab.bits() |
                                PageFlags::PageHead.bits() | PageFlags::PageTail.bits();
    }
}

// 当页是复合页并且是头部页时使用,
// 如果页不是复合页, 则是单页
pub(crate) struct PageHead {
    pub(crate) lru: RawListNode<Page>,
    pub(crate) order: usize,
}

// 当页是复合页并且是尾随页时使用,
// 用于标记头部页地址和该复合页的 order
pub(crate) struct PageComp {
    comp_head: usize,
    comp_order: usize,
}

// 当页是 slab 页时使用, 记录 slab 相关信息
pub(crate) struct PageSlab {
    freelist: usize,
    inuse: usize,
}

pub(crate) enum PageType {
    Head(PageHead),
    Comp(PageComp),
    Slab(PageSlab),
}

pub struct Page {
    flags: PageFlags,
    pub(crate) ty: PageType,
    // 记录该页被 map 的次数.
    mapcount: AtomicI32,
    // 记录该页被引用的次数, 如果引用次数不为 0, 则不能释放.
    refcount: AtomicU32,
}

impl Page {
    #[inline]
    // 设置单页或者头部页的 order, 并且设置 PageBuddy 标记, 表明该页属于伙伴系统
    pub(crate) fn set_page_order(&mut self, order: usize) {
        match &mut self.ty {
            PageType::Head(h) => h.order = order,
            _ => {
                unreachable!()
            }
        }
        self.flags.set(PageFlags::PageBuddy, true);
    }

    // 清除单页或者头部页的 order, 并且清除 PageBuddy 标记, 表明可以被其他系统使用
    #[inline]
    pub(crate) fn rmv_page_order(&mut self) {
        self.flags.set(PageFlags::PageBuddy, false);
        match &mut self.ty {
            PageType::Head(h) => h.order = 0,
            _ => {
                unreachable!()
            }
        }
    }

    // 获取单页或者头部页的 order
    #[inline]
    fn page_order(&self) -> usize {
        unsafe {
            match &self.ty {
                PageType::Head(h) => h.order,
                _ => {
                    unreachable!()
                }
            }
        }
    }

    // 这个函数检查伙伴页是否空闲并且是真的一个伙伴页
    //
    // 我们可以合并一个页面和它的伙伴, 如果:
    // (a) 伙伴不是空洞的(调用前检查!)
    // (b) 伙伴是伙伴系统的
    // (c) page 和伙伴页有相同的 order
    // (d) page 和伙伴页有相同的 zone
    //
    // 为了记录一个 page 是否在伙伴系统中, 我们设置 PageBuddy.
    // 并且通过 Zone lock 来保证 set, clear, test PageBuddy 的序列化.
    //
    // 为了记录 order, 使用 page_order.
    #[inline]
    fn page_is_buddy(&self, buddy: &Page, order: usize) -> bool {
        if buddy.flags.contains(PageFlags::PageBuddy) && buddy.page_order() == order {
            if buddy.flags.bits() & PageFlags::ZoneMask.bits()
                != self.flags.bits() & PageFlags::ZoneMask.bits()
            {
                return false;
            }
            assert_eq!(
                buddy.refcount.load(core::sync::atomic::Ordering::Relaxed),
                0
            );
            return true;
        }
        false
    }

    // 当页是复合页时, 判断页是否是头部页
    #[inline]
    fn page_is_head(&self) -> bool {
        self.flags.contains(PageFlags::PageHead)
    }

    // 当页是复合页时, 判断页是否是尾随页
    #[inline]
    fn page_is_tail(&self) -> bool {
        self.flags.contains(PageFlags::PageTail)
    }

    // 判断页是否是一个复合页
    #[inline]
    fn page_is_compound(&self) -> bool {
        self.page_is_head() || self.page_is_tail()
    }

    // 当页是复合页时, 返回该复合页对应的头部页
    #[inline]
    fn page_head(&self) -> &Page {
        assert!(self.page_is_compound());
        assert_ne!(self.page_is_head(), self.page_is_tail());

        if self.page_is_head() {
            self
        } else {
            unsafe {
                match &self.ty {
                    PageType::Comp(c) => unsafe {
                        let head = core::ptr::read_volatile(&c.comp_head);
                        (head as *const Page).as_ref().unwrap()
                    },
                    _ => {
                        unreachable!()
                    }
                }
            }
        }
    }

    #[inline]
    fn compound_or_single_order(&self) -> usize {
        if self.page_is_tail() {
            unsafe {
                match &self.ty {
                    PageType::Comp(c) => c.comp_order,
                    _ => {
                        unreachable!()
                    }
                }
            }
        } else {
            self.page_order()
        }
    }

    #[inline]
    fn page_expected_state(&self, flags: PageFlags) -> Result<()> {
        if unlikely(self.mapcount.load(core::sync::atomic::Ordering::Relaxed) != -1) {
            return Err(crate::error::PaError::EMAPCOUNT);
        }
        if unlikely(self.refcount.load(core::sync::atomic::Ordering::Relaxed) != 0) {
            return Err(crate::error::PaError::EREFCOUNT);
        }
        if unlikely(self.flags.bits() & flags.bits() != 0) {
            return Err(crate::error::PaError::ECHECK);
        }
        Ok(())
    }

    // 释放页时的检测, 包括 mapcount/refcount/flags 检查
    #[inline]
    pub(crate) fn free_pages_check(&self) -> Result<()> {
        self.page_expected_state(PageFlags::CheckAtFree)
    }

    // 当页是复合页时, 检查该复合页是否正确
    #[inline]
    fn free_tail_pages_check(&self, page: &Page) -> Result<()> {
        if unlikely(!page.page_is_tail()) {
            return Err(crate::error::PaError::ENOTAIL);
        }
        if unlikely(page.page_head() as *const Page as usize != self as *const Page as usize) {
            return Err(crate::error::PaError::ENOCOMP);
        }
        Ok(())
    }

    // 预备释放页, 进行页检查
    #[inline(always)]
    pub(crate) fn free_pages_prepare(&mut self, order: usize) -> Result<()> {
        assert!(!self.page_is_tail());

        if unlikely(order != 0) {
            let compound = self.page_is_compound();
            assert!(compound);
            assert_eq!(self.compound_or_single_order(), order);

            for i in 1..(1 << order) {
                let next;
                unsafe {
                    next = (self as *mut Page).as_mut().unwrap();
                }

                if compound {
                    self.free_tail_pages_check(next)?
                }
                next.free_pages_check()?;
                next.flags &= !PageFlags::CheckAtPrep;
            }
        }
        self.free_pages_check()?;
        self.flags &= !PageFlags::CheckAtPrep;
        Ok(())
    }

    // 为配对中的匹配伙伴和它们组成的组合 O(n + 1)页定位结构页
    #[inline]
    fn find_buddy_pfn(page_pfn: PfnAddr, order: usize) -> PfnAddr {
        PfnAddr::from(page_pfn.value() ^ (1 << order))
    }

    #[inline]
    pub(crate) fn prefetch_buddy(&self) {
        let pfn = self.page_to_pfn();
        let buddy_pfn = Page::find_buddy_pfn(pfn, 0);
        let buddy;
        unsafe {
            buddy = (self as *const Page)
                .add(buddy_pfn.value() - pfn.value())
                .as_ref()
                .unwrap();
        }
        prefetch(buddy as *const Page);
    }

    #[inline(always)]
    pub(crate) fn buddy_list_del(&mut self) {
        unsafe {
            match &mut self.ty {
                PageType::Head(h) => h.lru.list_del(),
                _ => {
                    unreachable!()
                }
            }
        }
    }

    #[inline(always)]
    pub(crate) fn buddy_list_add(&mut self, head: &mut RawListHead<Page>) {
        match &mut self.ty {
            PageType::Head(h) => unsafe { head.list_add(&mut h.lru) },
            _ => {
                unreachable!()
            }
        }
    }

    pub(crate) fn buddy_list_add_tail(&mut self, head: &mut RawListHead<Page>) {
        match &mut self.ty {
            PageType::Head(h) => unsafe { head.list_add_tail(&mut h.lru) },
            _ => {
                unreachable!()
            }
        }
    }

    #[inline]
    fn check_new_page(&self) -> Result<()> {
        self.page_expected_state(PageFlags::CheckAtPrep)
    }

    #[inline]
    pub(crate) fn check_new_pages(&self, order: usize) -> Result<()> {
        for i in 0..1 << order {
            unsafe {
                let p = (self as *const Page).add(i).as_ref().unwrap();
                p.check_new_page()?
            }
        }
        Ok(())
    }

    #[inline(always)]
    pub(crate) fn check_pcp_refill(&self) -> bool {
        self.check_new_page().is_ok()
    }

    #[inline(always)]
    pub(crate) fn check_new_pcp(&self) -> bool {
        self.check_new_page().is_ok()
    }

    #[inline]
    fn set_page_refcounted(&mut self) {
        assert!(!self.page_is_tail());
        assert!(self.refcount.load(core::sync::atomic::Ordering::Relaxed) == 0);
        self.refcount
            .store(1, core::sync::atomic::Ordering::Release);
    }

    #[inline]
    pub(crate) fn post_alloc_hook(&mut self, order: usize, gfp_mask: GfpFlags) {
        match &mut self.ty {
            PageType::Head(h) => h.order = 0,
            _ => unreachable!(),
        }
        self.set_page_refcounted();
    }

    #[inline]
    fn set_compund_order(&mut self, order: usize) {
        unsafe {
            let next = (self as *mut Page).add(1).as_mut().unwrap();
            match &mut next.ty {
                PageType::Comp(c) => c.comp_order = order,
                _ => {
                    unreachable!()
                }
            }
        }
    }

    #[inline]
    fn set_compund_head(&mut self, head: *mut Page) {
        match &mut self.ty {
            PageType::Comp(c) => unsafe {
                core::ptr::write_volatile(&mut c.comp_head, head as usize)
            },
            _ => {
                unreachable!()
            }
        }
    }

    #[inline]
    pub(crate) fn prep_compound_page(&mut self, order: usize) {
        let nr_pages = 1 << order;

        self.set_compund_order(order);
        self.flags.set(PageFlags::PageHead, true);
        for i in 1..nr_pages {
            unsafe {
                let p = (self as *mut Page).add(i).as_mut().unwrap();
                p.refcount.store(0, core::sync::atomic::Ordering::Release);
                p.set_compund_head(self as *mut Page);
            }
        }
    }

    #[inline]
    pub(crate) fn prep_new_page(&mut self, order: usize, gfp_mask: GfpFlags) {
        self.post_alloc_hook(order, gfp_mask);

        if gfp_mask.contains(GfpFlags::ZERO) {
            for i in 0..1 << order {
                unsafe {
                    let page = (self as *const Page).add(i).as_ref().unwrap();
                    let v = page.page_to_virtlm();
                    clear_page(v.value() as *mut c_void);
                }
            }
        }

        if order != 0 {
            self.prep_compound_page(order);
        }
    }

    #[inline]
    pub(crate) fn __free_one_page(
        &mut self,
        zone_inner: &mut ZoneInner,
        mut pfn: PfnAddr,
        mut order: usize,
    ) {
        assert!(!self.flags.contains(PageFlags::CheckAtPrep));
        assert!(pfn.value() & ((1 << order) - 1) == 0);

        let mut combined_pfn;
        let mut buddy_pfn = PfnAddr::from(0);
        let mut buddy;
        let mut page = self;
        while order < MAX_ORDER - 1 {
            buddy_pfn = Page::find_buddy_pfn(pfn, order);
            unsafe {
                buddy = (page as *mut Page).add(buddy_pfn.value() - pfn.value());
            }
            if !buddy_pfn.pfn_valid() {
                break;
            }
            unsafe {
                if !page.page_is_buddy(buddy.as_ref().unwrap(), order) {
                    break;
                }
                buddy.as_mut().unwrap().buddy_list_del();
            }
            zone_inner.free_area[order].nr_free -= 1;
            unsafe {
                buddy.as_mut().unwrap().rmv_page_order();
            }
            combined_pfn = PfnAddr::from(buddy_pfn.value() & pfn.value());
            unsafe {
                page = (page as *mut Page)
                    .add(combined_pfn.value() - pfn.value())
                    .as_mut()
                    .unwrap();
            }
            pfn = combined_pfn;
            order += 1;
        }

        page.set_page_order(order);

        if (order < MAX_ORDER - 2) && buddy_pfn.pfn_valid() {
            let mut higher_page;
            let mut higher_buddy;
            combined_pfn = PfnAddr::from(buddy_pfn.value() & pfn.value());
            unsafe {
                higher_page = (page as *mut Page).add(combined_pfn.value() - pfn.value());
                buddy_pfn = Page::find_buddy_pfn(combined_pfn, order + 1);
                higher_buddy = higher_page.add(buddy_pfn.value() - combined_pfn.value());
            }
            unsafe {
                if buddy_pfn.pfn_valid()
                    && higher_page
                        .as_ref()
                        .unwrap()
                        .page_is_buddy(higher_buddy.as_ref().unwrap(), order + 1)
                {
                    page.buddy_list_add(&mut zone_inner.free_area[order].free_list);
                    zone_inner.free_area[order].nr_free += 1;
                }
            }
        }
        page.buddy_list_add(&mut zone_inner.free_area[order].free_list);
        zone_inner.free_area[order].nr_free += 1;
    }

    pub(crate) fn free_pcp_prepare(&mut self) -> Result<()> {
        self.free_pages_prepare(0)
    }

    fn free_unref_page_commit(&mut self) {
        let zone = if self.flags.contains(PageFlags::ZoneNormal) {
            zone_idx(crate::zone::ZoneType::Normal)
        } else {
            zone_idx(crate::zone::ZoneType::Dma)
        };

        let raw_pcp;
        let mut pcp;
        unsafe {
            raw_pcp =
                zone.inner.lock_irq_save().pcp.get_mut(this_processor_id()) as *mut PerdataPages;
            pcp = raw_pcp.as_mut().unwrap();
        }
        self.buddy_list_add(&mut pcp.lists);
        pcp.count += 1;
        if pcp.count >= pcp.high {
            let batch;
            unsafe {
                batch = core::ptr::read_volatile(&pcp.batch);
            }
            zone.free_pcppages_bulk(batch as usize, pcp);
        }
    }

    pub(crate) fn free_unref_page(&mut self) {
        self.free_pcp_prepare().unwrap();
        let flags = local_irq_save();
        self.free_unref_page_commit();
        local_irq_restore(flags);
    }

    pub(crate) fn free_pages_ok(&mut self, order: usize) {
        self.free_pages_prepare(order).unwrap();
        let zone = if self.flags.contains(PageFlags::ZoneNormal) {
            zone_idx(crate::zone::ZoneType::Normal)
        } else {
            zone_idx(crate::zone::ZoneType::Dma)
        };
        let flags = local_irq_save();
        zone.free_one_page(self, self.page_to_pfn(), order);
        local_irq_restore(flags);
    }
}

extern "C" {
    pub fn clear_page(to: *mut core::ffi::c_void);
}
